/*
 * ==================================
 * =========== PROJECT 4 ============
 * == Partial Reliability Protocol ==
 * ==================================
 *
 * Name:            Luke Valenta
 * Directory ID:    lvalenta
 * Student ID:      112113230
 */

 /*
  * The sender implementation of the Partial Reliability Protocol
  * See the README file for more details.
 */

#include <stdio.h>
#include <ctype.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <sys/time.h>
#include <time.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#define MAX_RTX 8 
#define TIMEOUT 100
#define MAX_PACKETS 65535
#define MESSAGE_SIZE 512
 
#define DROP_RATE 15
#define DEBUG 0

typedef enum {
    false,
    true
} bool;

typedef enum {
    TX,
    RTX
} event;

typedef struct {
    struct timeval send_time;
    unsigned sn;
    int rtx_count;
    bool sent;
    bool acked;
} packet;

/* Globals */
char * g_ip;                    // receiver's IP address
char * g_log_file;              // name of the log file
FILE * g_log_fp;                // log file
unsigned g_port;                // receiver's port
unsigned g_num;                 // total number of packets to transmit
unsigned g_win;                 // sender window size
unsigned g_loss;                // acceptible loss rate per window
struct timeval g_start_time;    // time of program start

/* Sliding window globals */
packet g_packets[MAX_PACKETS];  // The packets currently in the window
unsigned g_base;                // The first sequence number in the current window
bool g_timeout;                 // An indicator of whether or not a timeout has occured
bool g_done;                    // An indicator of whether or not the final acks have been received

/* Socket globals */
struct sockaddr_in g_recvaddr;
int g_sockfd;

/* */
bool dropOrNot();

/* Function prototypes */
int parse_arg(int argc, char ** argv);
void usage(char *err_msg, char *name);
void setup(int argc, char ** argv);
void cleanup();
void log_send(unsigned sn, event e);
void log_shift(unsigned old_start, unsigned new_start, int * bitmap, int losses);
void log_timeout();
int get_elapsed_time_in_msec(struct timeval start);
int timeval_subtract (struct timeval * result, struct timeval * x, struct timeval * y);
int get_socket();
void construct_address(struct sockaddr_in * recvaddr, char * ip, int port);
void t_send(unsigned sn);
void init_window();
void transmit();
void process_ack(unsigned sn, int * bitmap);
void shift_window();

int parse_arg(int argc, char ** argv) {
    char c;
    bool got_ip = false;
    bool got_port = false;
    bool got_loss = false;
    bool got_num = false;
    g_ip = NULL;
    g_log_file = "stdout";
    g_log_fp = stdout;
    g_port = -1;
    g_num = 10;
    g_win = 128;
    g_loss = -1;
    while((c = getopt(argc, argv, "i:p:l:n:w:f:")) != -1) {
        switch (c)
        {
            case 'i':
               g_ip = malloc(strlen(optarg) + 1);
               strcpy(g_ip, optarg);
               got_ip = true;
               break;
            case 'p':
               g_port = atoi(optarg);
               got_port = true;
               break;
            case 'l':
               g_loss = atoi(optarg);
               got_loss = true;
               break;
            case 'n':
               got_num = true;
               g_num = atoi(optarg);
               break;
            case 'w':
               g_win = atoi(optarg);
               break;
            case 'f':
               g_log_file = malloc(strlen(optarg) + 1);
               strcpy(g_log_file, optarg);
               break;
            case '?':
               usage ("", argv[0]);
               break;
        }
    }
    if (optind != argc) {
        usage("", argv[0]);
    }
    if (!got_ip || !got_port || !got_loss || !got_num) {
        usage("", argv[0]);
    }

    return optind;
}

void usage (char *err_msg, char *name) {
    fprintf(stderr, "\n%s\nUsage: %s -i <IP> -p <PORT> -l <LOSS_RATE> -n <NUM> [-w <WIN>] [-f <LOG_FILE>]\n",
            err_msg, name);
    exit(1);
}

void log_send(unsigned sn, event e) {
    fprintf(g_log_fp, "TIME: %dms, \tSN: %u, \tEVENT: %s\n\n", get_elapsed_time_in_msec(g_start_time), sn, (e == TX) ? "TX" : "RTX");
}

void log_shift(unsigned old_start, unsigned new_start, int *bitmap, int losses) {
    int i;
    fprintf(g_log_fp, "TIME: %dms, \tOLD_START: %u, \tNEW_START: %u, \tLOSSES: %d\n",
            get_elapsed_time_in_msec(g_start_time), old_start, new_start, losses);
    fprintf(g_log_fp, "WIN_BITMAP:");
    for(i = 0; i < g_win; ++i) {
        if(i % 64 == 0) {
            fprintf(g_log_fp, "\n\t");
        }
        fprintf(g_log_fp, "%d", bitmap[i]);
        if(i % 8 == 7) {
            fprintf(g_log_fp, " ");
        }
    }
    fprintf(g_log_fp, "\n\n");

}

void log_timeout() {
    fprintf(g_log_fp, "TIME: %dms, \tRECEIVER TIMEOUT\n\n", get_elapsed_time_in_msec(g_start_time));
}

int get_elapsed_time_in_msec(struct timeval start) {
    struct timeval start_time, curr_time, diff_time;
    start_time.tv_sec = start.tv_sec;
    start_time.tv_usec = start.tv_usec;
    gettimeofday(&curr_time, NULL);
    timeval_subtract(&diff_time, &curr_time, &start_time);
    return diff_time.tv_sec*1000 + diff_time.tv_usec/1000;
}

/* http://www.gnu.org/software/libc/manual/html_node/Elapsed-Time.html */
int timeval_subtract (struct timeval *result, struct timeval *x, struct timeval *y) {
    /* Perform the carry for the later subtraction by updating y. */
    if (x->tv_usec < y->tv_usec) {
        int nsec = (y->tv_usec - x->tv_usec) / 1000000 + 1;
        y->tv_usec -= 1000000 * nsec;
        y->tv_sec += nsec;
    }
    if (x->tv_usec - y->tv_usec > 1000000) {
        int nsec = (x->tv_usec - y->tv_usec) / 1000000;
        y->tv_usec += 1000000 * nsec;
        y->tv_sec -= nsec;
    }

    /* Compute the time remaining to wait.
     *           tv_usec is certainly positive. */
    result->tv_sec = x->tv_sec - y->tv_sec;
    result->tv_usec = x->tv_usec - y->tv_usec;

    /* Return 1 if result is negative. */
    return x->tv_sec < y->tv_sec;
}

/* Attempt to transmit a packet */
void transmit() {
    int sn;
    // Look through packets in window
    for (sn = g_base; sn < g_base + g_win && sn < g_num; ++sn) {
        if (g_packets[sn].sent == false) { // Not yet sent
            gettimeofday(&g_packets[sn].send_time, NULL);
            g_packets[sn].sent = true;
            log_send(sn, TX);
            t_send(sn);
            return;
        } else if (g_packets[sn].acked == false) { // Sent but not acked
            // Check if we should retransmit this pakcet (timed out)
            if(get_elapsed_time_in_msec(g_packets[sn].send_time) > TIMEOUT) {
                if(g_packets[sn].rtx_count >= MAX_RTX) { // RECEIVER TIMEOUT
                    log_timeout();
                    g_timeout = true;
                    return;
                } else { // retransmit packet
                    gettimeofday(&g_packets[sn].send_time, NULL);
                    g_packets[sn].rtx_count++;
                    log_send(sn, RTX);
                    t_send(sn);
                    return;
                }
            }
        }
    }
    // No packets are ready to be sent... wait for a timeout or an ack
}

void check_for_acks() {
    int recvlen = 0;
    int buflen = 2 + g_win/8 + 1; // Size of ack message
    unsigned base;
    int i;
    int ret;
    int * bitmap;
    char * buffer;
    struct timeval selTimeout;       /* Timeout for select call */
    fd_set sockSet;                  /* Set of sockets to pass to select */

    /* Set the select timeout */
    selTimeout.tv_sec = 0;
    selTimeout.tv_usec = 0;

    buffer = malloc(buflen);
    bitmap = malloc(sizeof(int) * g_win);
    
    // Check if there is a waiting ack
    FD_ZERO(&sockSet);
    FD_SET(g_sockfd, &sockSet);
    if ((ret = select (g_sockfd + 1, &sockSet, NULL, NULL, &selTimeout) > 0)) {
        if (FD_ISSET(g_sockfd, &sockSet)) {
            //if(DEBUG) fprintf(stderr, "DEBUG: received an ack from receiver\n");
            recvlen = recvfrom(g_sockfd, buffer, buflen, 0, 0, 0);
            if(recvlen < 0) {
                fprintf(stderr, "recvfrom() failed\n");
                exit(EXIT_FAILURE);
            }
            //if(DEBUG) fprintf(stderr, "DEBUG: recvlen: %d\n", recvlen);
        }
    }

    // Some acks were received
    if(ret != 0) {
        base = ((((unsigned)buffer[0]) & 0xFF) << 8) + (((unsigned)buffer[1]) & 0xFF);
        for(i = 0; i < g_win; ++i) {
            bitmap[i] = (buffer[2 + i/8] >> (7 - (i % 8))) & 0x1;
        }
        process_ack(base, bitmap);
    }

    free(bitmap);
    free(buffer);
}


void process_ack(unsigned base, int * bitmap) {
    unsigned i;
    for (i = 0; (i < g_win) && (base + i < g_base + g_win); ++i) {
        if(bitmap[i] == 1 && g_packets[base + i].acked == false) {
            if(DEBUG) fprintf(stderr, "DEBUG: received ack for packet %d\n", base + i);
            g_packets[base + i].acked = true;
        }
    }
    shift_window();
}

void shift_window() {
    int i;
    int num_lost = 0;
    int losses;
    int max_allowable_losses = (g_win * g_loss)/100;
    int old_start = g_base;
    int new_start = g_base;
    int * bitmap;

    // Count the total number of un-acked packets
    for(i = 0; i < g_win; ++i) {
        if(g_packets[old_start + i].acked == false) {
            num_lost++;
        }
    }
    // The number of losses in the current window
    losses = num_lost;

    // Move the start index up until the number of losses in the window reaches the max allowable
    while(losses <= max_allowable_losses) {
        if(g_packets[new_start++].acked == true) {
            losses++;
        }
    }
    g_base = new_start;

    // Create bitmap of send window
    bitmap = malloc(sizeof(int) * g_win);
    for(i = 0; i < g_win; ++i) {
        if(g_packets[new_start + i].acked == true) {
            bitmap[i] = 1;
        } else {
            bitmap[i] = 0;
        }
    }
    if(new_start != old_start) { // We have shifted the window
        log_shift(old_start, new_start, bitmap, num_lost);
    }

    // Stopping condition. Stop if the final window is full enough
    if(g_base >= g_num - g_win && num_lost <= max_allowable_losses) {
        g_done = true;
    }

    free(bitmap);
}

void t_send(unsigned sn) {
    char message[MESSAGE_SIZE];
    int i;

    // Drop a packet with a certain probability
    if (dropOrNot() == true) {
        //if(DEBUG) fprintf(stderr, "DEBUG: packet %u dropped\n\n", sn);
        return;
    }

    // Construct message
    message[0] = (sn >> 8) & 0xFF;
    message[1] = (sn) & 0xFF;
    for (i = 2; i < MESSAGE_SIZE; ++i) {
        message[i] = (char) rand();
    }

    // Send message to receiver
    if (sendto(g_sockfd, message, MESSAGE_SIZE, 0, (struct sockaddr *) &g_recvaddr, sizeof(struct sockaddr_in)) < 0) {
        printf("sendto() failed");
        exit(1);
    }
}

bool dropOrNot() {
    int coin = rand() % 100;
    if (coin > DROP_RATE) {
        return false;
    } 
    return true;
}

/* Construct receiver address */
void construct_address(struct sockaddr_in * recvaddr, char *ip, int port) {
    int ret;
    memset(recvaddr, 0, sizeof(struct sockaddr_in)); // Zero out fields
    recvaddr->sin_family = AF_INET; // IPv4 address family
    recvaddr->sin_port = htons(port); // Change the port to network byte order
    ret = inet_pton(AF_INET, ip, &recvaddr->sin_addr); // Convert address
    if (ret == 0) {
        printf("inet_pton() invalid address string");
        exit(1);
    } else if (ret < 0) {
        printf("inet_pton() failed");
        exit(1);
    }
}

/* Open a UDP socket */
int get_socket() {
    int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sockfd < 0) {
        printf("socket() failed");
        exit(1);
    }
    return sockfd;
}

/* Initialize the window */
void init_window() {
    int i;
    g_base = 0; // The first sequence number
    // Initialize all the packets that we will be sending
    for (i = 0; i < g_num; ++i) {
        //g_packets[i].send_time = {0, 0};
        g_packets[i].sn = g_base + i;
        g_packets[i].rtx_count = 0;
        g_packets[i].sent = false;
        g_packets[i].acked = false;
    }
}

void setup(int argc, char ** argv) {
    parse_arg(argc, argv);
    gettimeofday(&g_start_time, NULL);
    if(strcmp(g_log_file, "stdout") != 0) {
        g_log_fp = fopen(g_log_file, "w");
    }
    construct_address(&g_recvaddr, g_ip, g_port);
    g_sockfd = get_socket();
    g_timeout = false;
    g_done = false;
    init_window();
}

void cleanup() {
    if(strcmp(g_log_file, "stdout") != 0) {
        fclose(g_log_fp);
        free(g_log_file);
    }
    free(g_ip);
}

int main(int argc, char ** argv) {

    setup(argc, argv);

    /* Main loop of program */
    while(g_done == false && g_timeout == false) {
        // Send a single packet, if one is ready to be sent
        transmit();
        
        // Check for a waiting ack
        check_for_acks();

    } // End while

    cleanup();
    return 0;
}
